import argparse
from pathlib import Path
import logging
from ultralytics import YOLO
from ultralytics import settings
import os
from utils import setup_logging, get_device_info, log_training_params, load_yaml_config, merge_configs
from utils.dataset import get_dataset_info
from utils.utils import copy_checkpoint_models, rename_log_file
from utils.paths import YOLO_SERVICE_DIR, PRETRAINED_MODELS_DIR, CHECKPOINTS_DIR, CONFIGS_DIR, RUNS_DIR # , setup_directories

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# 初始化日志记录器
logger = logging.getLogger("YOLO_Training")

# 创建必要的目录结构
# setup_directories()

# 配置 Ultralytics 设置，指定权重保存目录
settings.update({'weights_dir': str(PRETRAINED_MODELS_DIR)})


def parse_args():
    parser = argparse.ArgumentParser(description="基于YOLO的交通工具检测训练")
    # YOLO 核心参数
    parser.add_argument("--data", type=str, default="./data.yaml", help="数据集路径")
    parser.add_argument("--epochs", type=int, default=32, help="训练轮数")
    parser.add_argument("--imgsz", type=int, default=640, help="图像尺寸")
    parser.add_argument("--batch", type=int, default=4, help="批次大小")
    parser.add_argument("--device", type=str, default='cpu', help="计算设备")
    parser.add_argument("--workers", type=int, default=None, help="数据加载线程数")
    parser.add_argument("--project", type=str, default=None, help="输出目录")
    parser.add_argument("--name", type=str, default=None, help="运行名称")
    parser.add_argument("--exist_ok", action="store_true", help="是否覆盖现有目录")
    parser.add_argument("--weights", type=str, default="yolov8m.pt", help="预训练模型")
    # 可选高频 YOLO 参数
    parser.add_argument("--lr0", type=float, default=None, help="初始学习率")
    parser.add_argument("--optimizer", type=str, default=None, help="优化器类型 (如 SGD, Adam)")
    # 项目自定义参数
    parser.add_argument("--log_encoding", type=str, default="utf-8-sig", help="日志编码格式")
    parser.add_argument("--use_yaml", type=bool, default=True, help="是否使用 YAML 配置")
    parser.add_argument("--log_level", type=str, default="INFO", help="日志级别")
    return parser.parse_args()


# @log_execution_time(logger_name="YOLO_Training")
def train_model(model, yolo_args, logger):
    """
    执行 YOLO 模型训练。

    Args:
        model (YOLO): YOLO 模型实例。
        yolo_args (argparse.Namespace): YOLO 训练参数。
        logger (logging.Logger): 日志记录器。

    Returns:
        results: 训练结果。
    """
    results = model.train(**vars(yolo_args))
    return results


def main():
    args = parse_args()
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    model_name = args.weights.replace('.pt', '')  # 动态获取模型名
    logger = setup_logging(
         base_path=YOLO_SERVICE_DIR,
         log_type="train",
         model_name=model_name,
         encoding=args.log_encoding,
         log_level=log_level,
        temp_log=True
     )
    logger.info("===== YOLOv8 交通工具检测训练开始 =====")
    try:
        yaml_config = {}
        if args.use_yaml:
            yaml_config = load_yaml_config(config_path=CONFIGS_DIR / "train.yaml", config_type="train")
        yolo_args, project_args = merge_configs(args, yaml_config)

        model_path = PRETRAINED_MODELS_DIR / project_args.weights
        log_training_params(project_args, model_path, CONFIGS_DIR, logger)

        logger.info(f"加载模型: {model_path}")
        if not model_path.exists():
            logger.error(f"模型文件不存在: {model_path}")
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        model = YOLO(str(model_path))

        if not Path(project_args.data).exists():
            logger.error(f"数据集配置文件不存在: {project_args.data}")
            raise FileNotFoundError(f"数据集配置文件不存在: {project_args.data}")
        else:
            logger.info(f"数据集配置文件存在: {project_args.data}")

        # device = yolo_args.device if hasattr(yolo_args, 'device') else ("0" if torch.cuda.is_available() else "cpu")
        # logger.info(f"使用设备: {device}")

        # logger.info(f"开始训练: epochs={yolo_args.epochs}, imgsz={yolo_args.imgsz}, batch={yolo_args.batch}")
        # 核心训练
        # results = model.train(**vars(yolo_args))
        # 调用封装的训练函数
        results = train_model(model, yolo_args, logger)

        if results and hasattr(results, 'save_dir'):
            copy_checkpoint_models(results.save_dir, project_args.weights, logger)
            rename_log_file(logger, results.save_dir, model_name, args.log_encoding)
        else:
            logger.warning("训练未生成有效的目录,跳过模型复制")
        logger.info("训练完成")
        logger.info(f"模型保存至: {results.save_dir}")

        metrics = results.results_dict
        log_lines = [
            f"训练结果:",
            f"- 精确率: {metrics['metrics/precision(B)']:.4f}",
            f"- 召回率: {metrics['metrics/recall(B)']:.4f}",
            f"- mAP@0.5: {metrics['metrics/mAP50(B)']:.4f}",
            f"- mAP@0.5:0.95: {metrics['metrics/mAP50-95(B)']:.4f}"
        ]
        logger.info("\n".join(log_lines))
    except Exception as e:
        logger.error(f"训练失败: {e}", exc_info=True)
        raise
    finally:
        logger.info("===== YOLOv8 交通工具检测训练结束 =====")


if __name__ == '__main__':
    # 配置命令行参数解析器
    parser = argparse.ArgumentParser(description="YOLOv8 Safety Helmet Detection Training")
    parser.add_argument('--data', type=str, default='data.yaml', help='数据集配置文件名（如 data.yaml）')
    parser.add_argument('--epochs', type=int, default=32, help='训练轮数（默认: 4）')
    parser.add_argument('--imgsz', type=int, default=640, help='输入图像尺寸（默认: 640）')
    parser.add_argument('--batch', type=int, default=4, help='批次大小（默认: 16）')
    parser.add_argument('--device', type=str, default='cpu', help='计算设备，如 0（GPU）或 cpu（默认: 0）')
    parser.add_argument('--workers', type=int, default=8, help='数据加载线程数（默认: 8）')
    parser.add_argument('--weights', type=str, default='yolov8n.pt', help='预训练模型文件（如 yolov8n.pt）')
    parser.add_argument('--run-name', type=str, default='train', help='训练运行名称（默认: train）')
    parser.add_argument('--resume', action='store_true', help='从检查点恢复训练（默认: False）')
    parser.add_argument('--log-encoding', type=str, default='utf-8-sig', choices=['utf-8', 'utf-8-sig', 'gbk'],
                        help='日志文件编码，utf-8-sig 适合 Windows（默认: utf-8-sig）')
    parser.add_argument('--optimizer', type=str, default='AdamW', choices=['AdamW', 'SGD', 'Adam'],
                        help='优化器类型（默认: AdamW）')
    parser.add_argument('--lr0', type=float, default=0.001, help='初始学习率（默认: 0.001）')
    parser.add_argument('--lrf', type=float, default=0.01, help='最终学习率（默认: 0.01）')
    parser.add_argument('--weight-decay', type=float, default=0.0005, help='权重衰减（默认: 0.0005）')
    parser.add_argument('--patience', type=int, default=10, help='早停耐心值（默认: 50）')
    parser.add_argument('--cos-lr', action='store_true', default=True, help='使用余弦学习率调度（默认: True）')
    parser.add_argument('--save-period', type=int, default=100, help='模型保存间隔（默认: 100）')
    parser.add_argument('--degrees', type=float, default=15.0, help='图像旋转角度范围（默认: 15.0）')
    parser.add_argument('--translate', type=float, default=0.1, help='图像平移范围（默认: 0.1）')
    parser.add_argument('--scale', type=float, default=0.5, help='图像缩放范围（默认: 0.5）')
    parser.add_argument('--mosaic', type=float, default=1.0, help='Mosaic 数据增强概率（默认: 1.0)')

    # 解析命令行参数
    args = parser.parse_args()
    # 切换到项目根目录
    os.chdir(YOLO_SERVICE_DIR)
    logger.info(f"工作目录: {os.getcwd()}")
    # 调用主函数
    main()
